/**
 * @Author: dn-jinmin
 * @File:  routerchain
 * @Version: 1.0.0
 * @Date: 2024/5/15
 * @Description:
 */

package chainx

import (
	"context"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/memory"
	"github.com/tmc/langchaingo/outputparser"
	"github.com/tmc/langchaingo/prompts"
	"github.com/tmc/langchaingo/schema"
)

type Router struct {
	Destination string
	Key         string
	Prompt      prompts.FormatPrompter
	Chains      chains.Chain
}

type MultiPromptInfo struct {
	Key         string
	Destination string
	Prompt      prompts.FormatPrompter
}

// 多提示词路由
type MultiChains struct {
	LLMChain     *chains.LLMChain
	Routers      map[string]*Router
	Destinations map[string]string
	OutputParser outputparser.Structured
}

func NewMultiChains(llm llms.Model, multiPrompts []*MultiPromptInfo) *MultiChains {
	var (
		routerLlMChain = make(map[string]*Router, len(multiPrompts))
		destination    = make(map[string]string, len(multiPrompts))
	)

	for i, _ := range multiPrompts {
		routerLlMChain[multiPrompts[i].Key] = &Router{
			Chains: chains.NewLLMChain(llm, multiPrompts[i].Prompt),
		}
		destination[multiPrompts[i].Key] = multiPrompts[i].Destination
	}

	return &MultiChains{
		Destinations: destination,
		Routers:      routerLlMChain,
		LLMChain:     chains.NewLLMChain(llm, _prompt),
		OutputParser: _outputParser,
	}
}

func (c *MultiChains) Call(ctx context.Context, inputs map[string]any,
	options ...chains.ChainCallOption) (map[string]any, error) {

	destinations := inputs[_input]

	result, err := chains.Call(ctx, c.LLMChain, map[string]any{
		_input:        destinations,
		_destinations: c.Destinations,
		_formatting:   c.OutputParser.GetFormatInstructions(),
	}, options...)
	if err != nil {
		return nil, err
	}

	text, ok := result[_text]
	if !ok {
		return nil, chains.ErrNotFound
	}

	return c.processLLMResult(ctx, text.(string), inputs)
}

// GetMemory gets the memory of the chain.
func (c *MultiChains) GetMemory() schema.Memory {
	return memory.NewSimple()
}

// GetInputKeys returns the input keys the chain expects.
func (c *MultiChains) GetInputKeys() []string {
	return nil
}

// GetOutputKeys returns the output keys the chain returns.
func (c *MultiChains) GetOutputKeys() []string {
	return nil
}

func (c *MultiChains) processLLMResult(ctx context.Context, text string,
	inputs map[string]interface{}) (map[string]interface{},
	error) {
	data, err := c.OutputParser.Parse(text)
	if err != nil {
		return nil, err
	}
	next := data.(map[string]string)
	router := c.Routers[next[_destinations]].Chains

	return chains.Call(ctx, router, inputs)
}
